import sys
import cv2
import os
from pyzbar import pyzbar
from datetime import datetime
from PySide6.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, 
                             QHBoxLayout, QPushButton, QLabel, QFileDialog, 
                             QTableWidget, QTableWidgetItem, QHeaderView,
                             QSplitter, QFrame, QMessageBox, 
                             QLineEdit, QFormLayout, QGroupBox)
from PySide6.QtCore import QThread, Signal, Qt, QUrl
from PySide6.QtGui import QImage, QPixmap, QColor, QDesktopServices

class VideoProcessingThread(QThread):
    """视频处理线程，负责识别条形码和记录时间段"""
    frame_ready = Signal(QImage)
    barcode_updated = Signal(dict)
    segment_complete = Signal(dict)
    process_finished = Signal()
    progress_updated = Signal(int)
    error_occurred = Signal(str)

    def __init__(self, video_path, output_dir):
        super().__init__()
        self.video_path = video_path
        self.output_dir = output_dir
        self.running = False
        self.paused = False
        self.skip_frames = 5
        self.resize_scale = 0.7
        
        # 视频基础信息
        self.fps = 30
        self.total_frames = 0
        self.frame_width = 0
        self.frame_height = 0
        self.codec = cv2.VideoWriter_fourcc(*'mp4v')
        
        # 条形码跟踪数据
        self.tracking_segments = []  # 已完成的片段
        self.current_segment = None  # 当前跟踪的片段

    def run(self):
        try:
            self.running = True
            # 尝试打开视频文件
            cap = cv2.VideoCapture(self.video_path)
            
            # 检查视频是否成功打开
            if not cap.isOpened():
                self.error_occurred.emit(f"无法打开视频文件，请检查文件是否存在或格式是否支持：\n{self.video_path}")
                self.running = False
                return

            # 获取视频基础信息
            self.fps = cap.get(cv2.CAP_PROP_FPS) or 30
            self.total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            self.frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            self.frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            
            # 检查视频是否有效
            if self.total_frames <= 0:
                self.error_occurred.emit("无法读取视频内容，可能是文件损坏或格式不支持")
                cap.release()
                self.running = False
                return
            
            # 确保输出目录存在
            try:
                os.makedirs(self.output_dir, exist_ok=True)
            except Exception as e:
                self.error_occurred.emit(f"无法创建输出目录：{str(e)}")
                cap.release()
                self.running = False
                return

            frame_count = 0
            video_writer = None
            
            while self.running and frame_count < self.total_frames:
                if self.paused:
                    self.msleep(50)
                    continue
                    
                ret, frame = cap.read()
                if not ret:
                    break
                
                # 更新进度
                progress = int((frame_count / self.total_frames) * 100)
                self.progress_updated.emit(progress)
                
                # 跳帧处理
                process_frame = (frame_count % self.skip_frames == 0)
                
                # 降低分辨率用于显示
                display_frame = frame.copy()
                if self.resize_scale != 1.0:
                    display_frame = cv2.resize(
                        display_frame, (0, 0), 
                        fx=self.resize_scale, fy=self.resize_scale
                    )
                
                # 检测条形码
                detected_barcodes = []
                if process_frame:
                    detected_barcodes = self.detect_barcodes(frame, display_frame)
                    self.update_segments(detected_barcodes, frame_count)
                
                # 处理视频写入
                if self.current_segment and process_frame:
                    # 如果还没有创建写入器，创建它
                    if video_writer is None:
                        start_time = self.current_segment["start_time"]
                        filename = f"{start_time}_{self.current_segment['data']}.mp4"
                        output_path = os.path.join(self.output_dir, filename)
                        self.current_segment["output_path"] = output_path
                        
                        video_writer = cv2.VideoWriter(
                            output_path, self.codec, self.fps,
                            (self.frame_width, self.frame_height)
                        )
                        
                        # 检查写入器是否成功创建
                        if not video_writer.isOpened():
                            self.error_occurred.emit(f"无法创建输出视频文件：{output_path}")
                            cap.release()
                            self.running = False
                            return
                
                # 写入当前帧
                if self.current_segment and video_writer is not None:
                    video_writer.write(frame)
                elif not self.current_segment and video_writer is not None:
                    # 结束当前片段写入
                    video_writer.release()
                    video_writer = None
                    if self.tracking_segments and self.tracking_segments[-1]["data"] == self.current_segment["data"]:
                        self.segment_complete.emit(self.tracking_segments[-1])
                
                # 转换为QImage并发送到UI
                rgb_image = cv2.cvtColor(display_frame, cv2.COLOR_BGR2RGB)
                h, w, ch = rgb_image.shape
                bytes_per_line = ch * w
                q_image = QImage(rgb_image.data, w, h, bytes_per_line, QImage.Format_RGB888)
                self.frame_ready.emit(q_image)
                
                frame_count += 1
            
            # 处理最后一个片段
            if video_writer is not None:
                video_writer.release()
                if self.current_segment:
                    self.finalize_current_segment(frame_count)
                    self.segment_complete.emit(self.tracking_segments[-1])
            
            cap.release()
            self.running = False
            self.process_finished.emit()
            
        except Exception as e:
            self.error_occurred.emit(f"处理过程中发生错误：{str(e)}")
            self.running = False

    def detect_barcodes(self, frame, display_frame):
        """识别条形码并在显示帧上绘制"""
        detected_barcodes = []
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        barcodes = pyzbar.decode(gray)
        
        for barcode in barcodes:
            # 在原始帧上获取坐标，然后转换到显示帧
            (x, y, w, h) = barcode.rect
            scale = self.resize_scale
            
            # 在显示帧上绘制
            cv2.rectangle(
                display_frame, 
                (int(x*scale), int(y*scale)), 
                (int((x + w)*scale), int((y + h)*scale)), 
                (0, 255, 0), 2
            )
            
            barcode_data = barcode.data.decode("utf-8")
            barcode_type = barcode.type
            current_time = datetime.now().strftime("%Y%m%d_%H%M%S")
            
            # 在显示帧上添加文本
            text = f"{barcode_data} ({barcode_type})"
            cv2.putText(
                display_frame, text, 
                (int(x*scale), int(y*scale) - 10), 
                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2
            )
            
            detected_barcodes.append({
                "data": barcode_data,
                "type": barcode_type,
                "time": current_time
            })
        
        return detected_barcodes
    
    def update_segments(self, detected_barcodes, frame_count):
        """更新条形码片段信息"""
        if not detected_barcodes:
            # 如果没有检测到条形码且有当前片段，结束它
            if self.current_segment:
                self.finalize_current_segment(frame_count)
                self.segment_complete.emit(self.tracking_segments[-1])
            return
        
        # 取第一个检测到的条形码
        detected_code = detected_barcodes[0]
        
        # 情况1：没有当前跟踪的片段，开始新片段
        if not self.current_segment:
            self.start_new_segment(detected_code, frame_count)
        
        # 情况2：检测到的条形码与当前跟踪的不同，结束当前片段并开始新片段
        elif self.current_segment["data"] != detected_code["data"]:
            self.finalize_current_segment(frame_count)
            self.segment_complete.emit(self.tracking_segments[-1])
            self.start_new_segment(detected_code, frame_count)
        
        # 情况3：与当前跟踪的相同，更新结束帧
        else:
            self.current_segment["end_frame"] = frame_count
            self.current_segment["end_time"] = detected_code["time"]
            self.current_segment["duration"] = (
                (frame_count - self.current_segment["start_frame"]) / self.fps
            )
            self.barcode_updated.emit(self.current_segment)
    
    def start_new_segment(self, barcode_info, frame_count):
        """开始新的条形码片段跟踪"""
        self.current_segment = {
            "data": barcode_info["data"],
            "type": barcode_info["type"],
            "start_frame": frame_count,
            "end_frame": frame_count,
            "start_time": barcode_info["time"],
            "end_time": barcode_info["time"],
            "duration": 0,
            "output_path": ""
        }
        self.barcode_updated.emit(self.current_segment)
    
    def finalize_current_segment(self, frame_count):
        """结束当前片段跟踪并保存"""
        if self.current_segment:
            self.current_segment["end_frame"] = frame_count
            self.current_segment["end_time"] = datetime.now().strftime("%Y%m%d_%H%M%S")
            self.current_segment["duration"] = (
                (frame_count - self.current_segment["start_frame"]) / self.fps
            )
            self.tracking_segments.append(self.current_segment)
            self.current_segment = None

    def toggle_pause(self):
        self.paused = not self.paused
    
    def stop(self):
        self.running = False
        self.wait()

class BarcodeVideoSplitter(QMainWindow):
    """条形码视频分割器主窗口"""
    def __init__(self):
        super().__init__()
        self.setWindowTitle("条形码视频自动分割器")
        self.setGeometry(100, 100, 1200, 800)
        
        self.video_thread = None
        self.video_path = ""
        self.output_dir = os.path.join(os.getcwd(), "output_segments")
        
        self.init_ui()
    
    def init_ui(self):
        # 主部件和布局
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QVBoxLayout(central_widget)
        
        # 突出显示的导入视频按钮
        import_layout = QHBoxLayout()
        self.import_video_btn = QPushButton("📂 导入视频文件")
        self.import_video_btn.setStyleSheet("""
            QPushButton {
                background-color: #4CAF50;
                color: white;
                padding: 10px 20px;
                font-size: 14px;
                border-radius: 5px;
            }
            QPushButton:hover {
                background-color: #45a049;
            }
        """)
        self.import_video_btn.clicked.connect(self.browse_video)
        import_layout.addWidget(self.import_video_btn)
        import_layout.addStretch()
        
        # 配置区域
        config_group = QGroupBox("视频和输出设置")
        config_layout = QFormLayout()
        
        self.video_path_edit = QLineEdit()
        self.video_path_edit.setReadOnly(True)
        self.browse_video_btn = QPushButton("浏览...")
        self.browse_video_btn.clicked.connect(self.browse_video)
        
        self.output_dir_edit = QLineEdit(self.output_dir)
        self.browse_output_btn = QPushButton("浏览...")
        self.browse_output_btn.clicked.connect(self.browse_output_dir)
        
        video_layout = QHBoxLayout()
        video_layout.addWidget(self.video_path_edit)
        video_layout.addWidget(self.browse_video_btn)
        
        output_layout = QHBoxLayout()
        output_layout.addWidget(self.output_dir_edit)
        output_layout.addWidget(self.browse_output_btn)
        
        config_layout.addRow("视频文件路径:", video_layout)
        config_layout.addRow("输出目录路径:", output_layout)
        
        # 控制按钮
        control_layout = QHBoxLayout()
        
        self.process_btn = QPushButton("▶️ 开始处理")
        self.process_btn.setStyleSheet("""
            QPushButton {
                background-color: #2196F3;
                color: white;
                padding: 8px 16px;
                border-radius: 4px;
            }
            QPushButton:disabled {
                background-color: #cccccc;
                color: #666666;
            }
        """)
        self.process_btn.clicked.connect(self.start_processing)
        self.process_btn.setEnabled(False)
        
        self.pause_btn = QPushButton("⏸️ 暂停/继续")
        self.pause_btn.setEnabled(False)
        self.pause_btn.clicked.connect(self.toggle_pause)
        
        self.open_output_btn = QPushButton("📁 打开输出目录")
        self.open_output_btn.clicked.connect(self.open_output_dir)
        self.open_output_btn.setEnabled(False)
        
        self.skip_label = QLabel("跳帧数: 5")
        
        self.decrease_skip_btn = QPushButton("-")
        self.decrease_skip_btn.clicked.connect(lambda: self.change_skip_frames(-1))
        self.decrease_skip_btn.setEnabled(False)
        
        self.increase_skip_btn = QPushButton("+")
        self.increase_skip_btn.clicked.connect(lambda: self.change_skip_frames(1))
        self.increase_skip_btn.setEnabled(False)
        
        control_layout.addWidget(self.process_btn)
        control_layout.addWidget(self.pause_btn)
        control_layout.addWidget(self.skip_label)
        control_layout.addWidget(self.decrease_skip_btn)
        control_layout.addWidget(self.increase_skip_btn)
        control_layout.addStretch()
        control_layout.addWidget(self.open_output_btn)
        
        # 当前跟踪信息
        current_track_layout = QHBoxLayout()
        self.current_label = QLabel("当前跟踪: 无")
        self.time_range_label = QLabel("时间范围: --")
        current_track_layout.addWidget(self.current_label)
        current_track_layout.addWidget(self.time_range_label)
        current_track_layout.addStretch()
        
        # 中间视频和结果显示区域
        splitter = QSplitter(Qt.Horizontal)
        
        # 视频显示区域
        self.video_frame = QFrame()
        self.video_frame.setFrameShape(QFrame.StyledPanel)
        video_layout = QVBoxLayout(self.video_frame)
        
        self.video_label = QLabel("请点击上方的'导入视频文件'按钮选择视频")
        self.video_label.setAlignment(Qt.AlignCenter)
        self.video_label.setMinimumSize(400, 300)
        video_layout.addWidget(self.video_label)
        
        # 结果表格
        self.results_frame = QFrame()
        self.results_frame.setFrameShape(QFrame.StyledPanel)
        results_layout = QVBoxLayout(self.results_frame)
        
        self.results_label = QLabel("分割片段记录:")
        
        # 创建表格显示结果
        self.results_table = QTableWidget()
        self.results_table.setColumnCount(6)
        self.results_table.setHorizontalHeaderLabels(["序号", "快递单号", "类型", "开始时间", "结束时间", "保存路径"])
        
        # 设置表格自适应宽度
        header = self.results_table.horizontalHeader()
        header.setSectionResizeMode(0, QHeaderView.ResizeToContents)
        header.setSectionResizeMode(1, QHeaderView.ResizeToContents)
        header.setSectionResizeMode(2, QHeaderView.ResizeToContents)
        header.setSectionResizeMode(3, QHeaderView.ResizeToContents)
        header.setSectionResizeMode(4, QHeaderView.ResizeToContents)
        header.setSectionResizeMode(5, QHeaderView.Stretch)
        
        results_layout.addWidget(self.results_label)
        results_layout.addWidget(self.results_table)
        
        splitter.addWidget(self.video_frame)
        splitter.addWidget(self.results_frame)
        splitter.setSizes([600, 600])
        
        # 添加到主布局
        main_layout.addLayout(import_layout)  # 突出显示的导入按钮
        main_layout.addWidget(config_group)
        main_layout.addLayout(control_layout)
        main_layout.addLayout(current_track_layout)
        main_layout.addWidget(splitter)
    
    def browse_video(self):
        """选择视频文件，增加错误处理和反馈"""
        try:
            file_path, _ = QFileDialog.getOpenFileName(
                self, "选择视频文件", "", "视频文件 (*.mp4 *.avi *.mov *.mkv *.flv *.wmv)"
            )
            
            if file_path:
                # 检查文件是否存在
                if not os.path.exists(file_path):
                    QMessageBox.warning(self, "文件不存在", f"所选文件不存在：\n{file_path}")
                    return
                
                # 检查文件是否可读取
                if not os.access(file_path, os.R_OK):
                    QMessageBox.warning(self, "无法读取", f"没有权限读取该文件：\n{file_path}")
                    return
                
                # 尝试打开视频文件以验证格式
                cap = cv2.VideoCapture(file_path)
                if not cap.isOpened():
                    QMessageBox.warning(self, "格式不支持", f"无法打开该视频文件，可能格式不支持：\n{file_path}")
                    cap.release()
                    return
                
                # 检查是否有视频流
                ret, _ = cap.read()
                if not ret:
                    QMessageBox.warning(self, "无法读取内容", f"无法读取视频内容，可能文件损坏：\n{file_path}")
                    cap.release()
                    return
                
                cap.release()
                
                # 一切正常，更新路径
                self.video_path = file_path
                self.video_path_edit.setText(file_path)
                self.process_btn.setEnabled(True)
                self.statusBar().showMessage(f"已导入视频文件：{os.path.basename(file_path)}")
                
        except Exception as e:
            QMessageBox.critical(self, "错误", f"选择视频文件时发生错误：\n{str(e)}")
    
    def browse_output_dir(self):
        """选择输出目录"""
        try:
            dir_path = QFileDialog.getExistingDirectory(
                self, "选择输出目录", self.output_dir
            )
            
            if dir_path:
                # 检查目录是否可写
                if not os.access(dir_path, os.W_OK):
                    QMessageBox.warning(self, "无法写入", f"没有权限写入该目录：\n{dir_path}")
                    return
                
                self.output_dir = dir_path
                self.output_dir_edit.setText(dir_path)
        except Exception as e:
            QMessageBox.critical(self, "错误", f"选择输出目录时发生错误：\n{str(e)}")
    
    def open_output_dir(self):
        """打开输出目录"""
        if os.path.exists(self.output_dir):
            QDesktopServices.openUrl(QUrl.fromLocalFile(self.output_dir))
        else:
            QMessageBox.warning(self, "警告", "输出目录不存在!")
    
    def start_processing(self):
        """开始处理视频"""
        if not self.video_path or not os.path.exists(self.video_path):
            QMessageBox.warning(self, "警告", "请选择有效的视频文件!")
            return
        
        # 确保输出目录存在
        try:
            os.makedirs(self.output_dir, exist_ok=True)
        except Exception as e:
            QMessageBox.critical(self, "错误", f"无法创建输出目录：\n{str(e)}")
            return
        
        # 停止已有线程
        if self.video_thread and self.video_thread.isRunning():
            self.video_thread.stop()
        
        # 清空表格
        self.results_table.setRowCount(0)
        
        # 创建并启动处理线程
        self.video_thread = VideoProcessingThread(self.video_path, self.output_dir)
        self.video_thread.frame_ready.connect(self.update_video_frame)
        self.video_thread.barcode_updated.connect(self.update_current_tracking)
        self.video_thread.segment_complete.connect(self.add_segment_result)
        self.video_thread.process_finished.connect(self.process_complete)
        self.video_thread.progress_updated.connect(self.update_progress)
        self.video_thread.error_occurred.connect(self.show_error)
        self.video_thread.start()
        
        # 更新UI状态
        self.process_btn.setEnabled(False)
        self.pause_btn.setEnabled(True)
        self.decrease_skip_btn.setEnabled(True)
        self.increase_skip_btn.setEnabled(True)
        self.open_output_btn.setEnabled(True)
    
    def update_video_frame(self, q_image):
        """更新视频显示"""
        pixmap = QPixmap.fromImage(q_image)
        scaled_pixmap = pixmap.scaled(
            self.video_label.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation
        )
        self.video_label.setPixmap(scaled_pixmap)
    
    def update_current_tracking(self, segment_info):
        """更新当前跟踪信息"""
        self.current_label.setText(f"当前跟踪: {segment_info['data']} ({segment_info['type']})")
        self.time_range_label.setText(
            f"时间范围: {segment_info['start_time']} - {segment_info['end_time']} "
            f"(持续: {segment_info['duration']:.2f}秒)"
        )
    
    def add_segment_result(self, segment_info):
        """添加分割结果到表格"""
        row = self.results_table.rowCount()
        self.results_table.insertRow(row)
        
        # 序号
        self.results_table.setItem(row, 0, QTableWidgetItem(str(row + 1)))
        
        # 快递单号
        item = QTableWidgetItem(segment_info['data'])
        item.setFlags(item.flags() & ~Qt.ItemIsEditable)
        self.results_table.setItem(row, 1, item)
        
        # 类型
        item = QTableWidgetItem(segment_info['type'])
        item.setFlags(item.flags() & ~Qt.ItemIsEditable)
        self.results_table.setItem(row, 2, item)
        
        # 开始时间
        item = QTableWidgetItem(segment_info['start_time'])
        item.setFlags(item.flags() & ~Qt.ItemIsEditable)
        self.results_table.setItem(row, 3, item)
        
        # 结束时间
        item = QTableWidgetItem(segment_info['end_time'])
        item.setFlags(item.flags() & ~Qt.ItemIsEditable)
        self.results_table.setItem(row, 4, item)
        
        # 保存路径
        item = QTableWidgetItem(segment_info['output_path'])
        item.setFlags(item.flags() & ~Qt.ItemIsEditable)
        self.results_table.setItem(row, 5, item)
        
        # 隔行变色
        if row % 2 == 0:
            for col in range(6):
                self.results_table.item(row, col).setBackground(QColor(240, 240, 240))
    
    def toggle_pause(self):
        """暂停/继续处理"""
        if self.video_thread:
            self.video_thread.toggle_pause()
            status = "已暂停" if self.video_thread.paused else "处理中"
            self.statusBar().showMessage(status)
    
    def change_skip_frames(self, delta):
        """调整跳帧数"""
        if self.video_thread:
            new_skip = self.video_thread.skip_frames + delta
            if new_skip >= 1:
                self.video_thread.skip_frames = new_skip
                self.skip_label.setText(f"跳帧数: {self.video_thread.skip_frames}")
    
    def update_progress(self, value):
        """更新进度条"""
        self.statusBar().showMessage(f"处理进度: {value}%")
    
    def process_complete(self):
        """处理完成"""
        self.statusBar().showMessage("处理完成!")
        self.pause_btn.setEnabled(False)
        self.decrease_skip_btn.setEnabled(False)
        self.increase_skip_btn.setEnabled(False)
        self.process_btn.setEnabled(True)
        QMessageBox.information(self, "完成", f"视频分割完成，共生成 {self.results_table.rowCount()} 个片段!")
    
    def show_error(self, message):
        """显示错误信息"""
        QMessageBox.critical(self, "处理错误", message)
        # 重置UI状态
        self.statusBar().showMessage("处理已停止")
        self.pause_btn.setEnabled(False)
        self.decrease_skip_btn.setEnabled(False)
        self.increase_skip_btn.setEnabled(False)
        self.process_btn.setEnabled(True)

    def closeEvent(self, event):
        """关闭窗口时停止线程"""
        if self.video_thread and self.video_thread.isRunning():
            self.video_thread.stop()
        event.accept()

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = BarcodeVideoSplitter()
    window.show()
    sys.exit(app.exec())
